import numpy as np
import tensorflow as tf
import torch
import torch.nn.functional as F
bfloat16 = tf.bfloat16.as_numpy_dtype
data_type_emu = {bfloat16: 0, np.float16: 1, np.float32: 2}

def gen_golden_data_simple():
    
    # data_type = bfloat16
    data_type = np.float16
    # data_type = np.float32


    ## 核间均分，单核计算量对齐:
    # input_shape = [17, 1024]

    ## 核间均分，单核计算量非对齐:
    input_shape = [17, 1023]

    ## 核间不均分，单核计算量对齐:
    # input_shape = [8, 1024]

    ## 核间不均分，单核计算量非对齐:
    # input_shape = [8, 1023]

    input_p = np.random.uniform(0.01, 0.99, input_shape).astype(data_type)
    input_y = np.random.randint(0, 2, input_shape).astype(data_type)
    p_torch = torch.from_numpy(input_p.astype(np.float32))
    y_torch = torch.from_numpy(input_y.astype(np.float32))
    golden_torch = F.binary_cross_entropy(p_torch, y_torch, reduction='none')
    golden = golden_torch.numpy().astype(data_type)

    tiling = np.array([input_shape[0] * input_shape[1], data_type_emu[data_type]], dtype=np.uint32)
    tiling.tofile("./input/input_tiling.bin")
    input_p.astype(data_type).tofile("./input/input_p.bin")
    input_y.astype(data_type).tofile("./input/input_y.bin")
    golden.astype(data_type).tofile("./output/golden.bin")


if __name__ == "__main__":
    gen_golden_data_simple()
